
#include <iostream>
#include <ctype.h>

#include "opencv2/video/tracking.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/nonfree/features2d.hpp"
#include "opencv2/objdetect/objdetect.hpp"
#include "opencv2/imgproc/imgproc_c.h"
#include "opencv2/features2d/features2d.hpp"
#include "opencv2/calib3d/calib3d.hpp"

#include "ObjectDetector.hpp"

using namespace cv;
using namespace std;

// SIFT Matcher vars
vector<DMatch> enh_good_matches;
vector<KeyPoint> keypoints_img, keypoints_ref;
Mat descriptor_img, descriptor_ref;
int vote;
Mat ref_img;
Mat matched_pnts_ref, matched_pnts_img;

static void help()
{
	// print a welcome message, and the OpenCV version
	cout << "\nThis is a demo of Lukas-Kanade optical flow lkdemo(),\n"
		"Using OpenCV version %s\n" << CV_VERSION << "\n"
		<< endl;

	cout << "\nHot keys: \n"
		"\tESC - quit the program\n"
		"\tr - auto-initialize tracking\n"
		"\tc - delete all the points\n"
		"\tn - switch the \"night\" mode on/off\n"
		"To add/remove a feature point click it\n" << endl;
}

Point2f point;
bool addRemovePt = false;

static void onMouse( int event, int x, int y, int /*flags*/, void* /*param*/ )
{
	if( event == CV_EVENT_LBUTTONDOWN )
	{
		point = Point2f((float)x,(float)y);
		addRemovePt = true;
	}
}

int main( int argc, char** argv )
{
	
	VideoCapture cap;
	TermCriteria termcrit(CV_TERMCRIT_ITER|CV_TERMCRIT_EPS,20,0.03);
	Size subPixWinSize(10,10), winSize(31,31);

	const int MAX_COUNT = 500;
	bool needToInit = false;
	bool nightMode = false;

	if( argc == 1 || (argc == 2 && strlen(argv[1]) == 1 && isdigit(argv[1][0])))
		cap.open(argc == 2 ? argv[1][0] - '0' : 0);
	else if( argc == 2 )
		cap.open(argv[1]);

	if( !cap.isOpened() )
	{
		cout << "Could not initialize capturing...\n";
		return 0;
	}

	help();

	namedWindow( "LK Demo", 1 );
	setMouseCallback( "LK Demo", onMouse, 0 );

	Mat gray, prevGray, image;
	vector<Point2f> points[2];

	// read reference image
	ref_img = imread("Picture 49.jpg");

	// init SIFT
	SIFT SIFTdect = SIFT(400, 3, 0.05, 10, 1.6);
	SIFTdect(ref_img, noArray(), keypoints_ref, descriptor_ref, false);   // get descriptors on reference image

	for(;;)
	{
		Mat frame;
		cap >> frame;
		if( frame.empty() )
			break;

		frame.copyTo(image);
		cvtColor(image, gray, CV_BGR2GRAY);

		if( nightMode )
			image = Scalar::all(0);

		if( needToInit )
		{
			// automatic initialization
			// input: gray image, parameters e.g. 
			// output: points[1] (vector<Point2f>)
			// notes: no descriptor, just normal points
			// goodFeaturesToTrack(gray, points[1], MAX_COUNT, 0.01, 10, Mat(), 3, 0, 0.04);

			// from extracted points, get their corresponding sub-pixel corners
			// cornerSubPix(gray, points[1], subPixWinSize, Size(-1,-1), termcrit);
			
			// Perform initial match btw target and current		
			// Using 3rd version
			vote = callSIFTMatch(frame, descriptor_ref, keypoints_ref, ref_img,     // <-- used for webcam image
			     keypoints_img, descriptor_img, enh_good_matches);
 
			// Using original version
			//vote = callSIFTMatch(ref_img, frame);

			// using 2nd version
			//vote = callSIFTMatch(frame, descriptor_ref, keypoints_ref, ref_img);			

			// Feed the matched points to optical flow i.e. copy to points[1]
			XtractMatchingPoints(keypoints_ref, keypoints_img, enh_good_matches, matched_pnts_ref, matched_pnts_img);
			
			// get sub-pixel corners
			points[1].clear();
			for (int i=0; i<matched_pnts_img.cols; i++)
			{
				//points[1][i] = matched_pnts_img.at<Vec2f>(i);
				points[1].push_back(matched_pnts_img.at<Vec2f>(i));
			}
			
			cornerSubPix(gray, points[1], subPixWinSize, Size(-1,-1), termcrit);

			addRemovePt = false;
		}
		else if( !points[0].empty() )
		{
			vector<uchar> status;
			vector<float> err;

			if(prevGray.empty())
				gray.copyTo(prevGray);

			// get output matched points
			calcOpticalFlowPyrLK(prevGray, gray, points[0], points[1], status, err, winSize,
				3, termcrit, 0, 0.001);

			size_t i, k;
			for( i = k = 0; i < points[1].size(); i++ )
			{
				if( addRemovePt )
				{
					if( norm(point - points[1][i]) <= 5 )
					{
						addRemovePt = false;
						continue;
					}
				}

				if( !status[i] )
					continue;

				points[1][k++] = points[1][i];
				circle( image, points[1][i], 3, Scalar(0,255,0), -1, 8);
			}
			points[1].resize(k);
		}

		// SubPix corner processing
		if( addRemovePt && points[1].size() < (size_t)MAX_COUNT )
		{
			vector<Point2f> tmp;
			tmp.push_back(point);
			cornerSubPix( gray, tmp, winSize, cvSize(-1,-1), termcrit);
			points[1].push_back(tmp[0]);
			addRemovePt = false;
		}

		needToInit = false;
		imshow("LK Demo", image);

		char c = (char)waitKey(10);
		if( c == 27 )
			break;
		switch( c )
		{
		case 'r':
			needToInit = true;
			break;
		case 'c':
			points[1].clear();
			break;
		case 'n':
			nightMode = !nightMode;
			break;
		default:
			;
		}

		// copy current points/image to previous points/image
		std::swap(points[1], points[0]);
		swap(prevGray, gray);
	}

	return 0;
}
